#!/usr/bin/python3
# -*- coding: utf-8 -*-

import cv2
import numpy as np

np.set_printoptions(formatter={'float': '{: 0.4f}'.format})

import rospy
import message_filters

from geometry_msgs.msg import Transform,Pose,Point32

from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError


from bs_imu import bs_imu
from bs_imu import bs_gps_raw
from bs_img import bs_img_real
from bs_img import bs_img_base
from bs_img.bs_cfg_real_5m import *
from bs_img.bs_lie_solve import *
from bs_img.serial_pack import *

import time
'''
输入图像
输出相对的解算值
'''




class BS_Solver:
    def __init__(self):
        wa_img_topic = "/camera_wa/image_dh"

        self.img_bridge = CvBridge()

        self.invalid_frame_num_drogue = 0
        self.invalid_frame_max = 8     #if invaild frame < invalid_frame_max, output results of the last frame

        #Kalman Filter
        self.Q = np.eye(3)
        self.R = np.eye(3)
        self.x_drogue = np.array([np.nan, np.nan, np.nan])
        self.w_drogue = 1.5


        self.wa_img_sub = rospy.Subscriber(wa_img_topic, Image, self.wa_img_cb)

        #只去解算锥套
        self.wa_drogue_pub = rospy.Publisher("/bs_debuger/result/wa_drogue", Transform, queue_size=1)


    def wa_img_cb(self,img_msg):
        if img_msg != None:
            img_gray = self.img_bridge.imgmsg_to_cv2(img_msg, img_msg.encoding)
            _, img_bin = cv2.threshold(img_gray, 50, 0xff, cv2.THRESH_BINARY)
            # 如果相机倒置需要旋转一下
            # lf_img_gray = cv2.rotate(lf_img_gray, cv2.ROTATE_180)
            
            plane_t_vec = [np.nan]*3
            plane_quat = [np.nan]*4
            drogue_t_vec = [np.nan]*3
            
            gt_pts,_ = bs_img_real.img_to_pts(img_bin)
            
            # 求解锥套位置
            drogue_result_dict = bs_img_real.solve_drogue_wa(gt_pts, mapping_coeffs, distortion_center, stretch_matrix, taper_radius=r_drogue)
                
            drogue_t_vec = drogue_result_dict["t_vec"]
            print("drogue_t", drogue_t_vec)
            
            drogue_msg = Transform()
            drogue_msg.translation.x = drogue_t_vec[0]
            drogue_msg.translation.y = drogue_t_vec[1]
            drogue_msg.translation.z = drogue_t_vec[2]
            self.wa_drogue_pub.publish(drogue_msg)

            # KF for drogue
            drogue_is_visiable = True if not(np.all(np.isnan(drogue_t_vec))) else False
            if not(drogue_is_visiable):
                self.invalid_frame_num_drogue = self.invalid_frame_num_drogue + 1
                if self.invalid_frame_num_drogue <= self.invalid_frame_max:
                    self.drogue_res = self.x_drogue
                    if not(np.all(np.isnan(self.drogue_res))):
                        drogue_res_frame = serial_frame("drogue", True, round((self.drogue_res[2])*100),round(self.drogue_res[0]*100),round((self.drogue_res[1])*100))  
                    else:
                        drogue_res_frame = serial_frame("drogue",False,round(0),round(0),round(0))
                else:
                    drogue_res_frame = serial_frame("drogue",False,round(0),round(0),round(0))
            
            else:
                self.invalid_frame_num_drogue = 0

                if np.all(np.isnan(self.x_drogue)):   #the first frame
                    self.x_drogue = drogue_t_vec
                else:
                    miu = drogue_t_vec-self.x_drogue
                    self.x_drogue = self.x_drogue + np.matmul(self.Q,np.matmul(np.linalg.pinv(self.Q+self.R/self.w_drogue),miu.T))
                    self.w_drogue = 1.5/(1+np.matmul(miu,np.matmul(np.linalg.pinv(self.R),miu.T)))
                drogue_res_frame = serial_frame("drogue",True,round((self.x_drogue[2])*100),round(self.x_drogue[0]*100),round((self.x_drogue[1])*100))  

            uart.write(drogue_res_frame)
            time.sleep(0.0005)

if __name__ == "__main__":
    rospy.init_node('bs_solver_node', anonymous=True)
    port, baud =  '/dev/ttyUSB0', 115200
    uart = serial.Serial(port,baud,parity=serial.PARITY_NONE)
    bs_solver = BS_Solver()
    rospy.spin()
